// @HEADER
//*********************************************************************//
//  SiYuan: A numerical PDE solver                                     //
//  Copyright (2022) YUAN Xi                                           //
//  This Software is released under the BSD 2-Clause license detailed  //
//  in the file "LICENSE" in the top-level SiYuan directory            //
//*********************************************************************//
// @HEADER

#ifndef _EquationSet_Magnetostatics_impl_hpp_
#define _EquationSet_Magnetostatics_impl_hpp_

#include "Teuchos_Tuple.hpp"
#include "Teuchos_ParameterList.hpp"
#include "Teuchos_RCP.hpp"
#include "Teuchos_Assert.hpp"
#include "Teuchos_StandardParameterEntryValidators.hpp"

#include "Phalanx_DataLayout_MDALayout.hpp"
#include "Phalanx_FieldManager.hpp"

#include "Intrepid2_FunctionSpaceTools.hpp"

#include "Panzer_String_Utilities.hpp"
#include "Panzer_IntegrationRule.hpp"
#include "Panzer_BasisIRLayout.hpp"

// include evaluators here
#include "Panzer_Integrator_BasisTimesVector.hpp"
#include "Panzer_Integrator_CurlBasisDotVector.hpp"
#include "Panzer_ScalarToVector.hpp"
#include "Panzer_Sum.hpp"
#include "Panzer_Constant.hpp"
#include "Panzer_BlockedDOFManager.hpp"
#include "Panzer_BlockedTpetraLinearObjFactory.hpp"
#include "Panzer_BlockedEpetraLinearObjFactory.hpp"
#include "Panzer_ReorderADValues_Evaluator.hpp"
#include "Panzer_LOCPair_GlobalEvaluationData.hpp"
#include "Panzer_CommonArrayFactories.hpp"

// ***********************************************************************
template <typename EvalT>
SiYuan::EquationSet_Magnetostatics<EvalT>::
EquationSet_Magnetostatics(const Teuchos::RCP<Teuchos::ParameterList>& params,
                   const int& default_integration_order,
                   const panzer::CellData& cell_data,
                   const Teuchos::RCP<panzer::GlobalData>& global_data,
                   const bool build_transient_support) :
  EquationSet<EvalT>(params,default_integration_order,cell_data,global_data,build_transient_support )
{
  // ********************
  // Validate and parse parameter list
  // ********************
  {    
    Teuchos::ParameterList valid_parameters;
    this->setDefaultValidParameters(valid_parameters);

    valid_parameters.set("Model ID","","Closure model id associated with this equaiton set");
    valid_parameters.set("DOF Name","VectorPotential","DOF Name");
    valid_parameters.set("Material Name","?","Material Name");
    valid_parameters.set("Basis Type","HCurl","Type of Basis to use");
    valid_parameters.set("Basis Order",1,"Order of the basis");
    valid_parameters.set("Integration Order",-1,"Order of the integration rule");
    valid_parameters.set< Teuchos::Array<double> >("Magnetization", Teuchos::tuple<double>(0.0),"Imposed Magnets");
    valid_parameters.set< Teuchos::Array<double> >("Current Density", Teuchos::tuple<double>(0.0),"Current Density");

    params->validateParametersAndSetDefaults(valid_parameters);
  }

  std::string model_id = params->get<std::string>("Model ID");
  m_dof_name = params->get<std::string>("DOF Name");
  m_matl_name = params->get<std::string>("Material Name");
  std::string basis_type = params->get<std::string>("Basis Type");
  int basis_order = params->get<int>("Basis Order");
  int integration_order = params->get<int>("Integration Order");

  // ********************
  // Setup DOFs and closure models
  // ********************
  {
    this->addDOF(m_dof_name,basis_type,basis_order,integration_order);
    this->addDOFCurl(m_dof_name);
  }

  this->addClosureModel(model_id);
  this->setupDOFs();

  this->addC0("Inverse of Permeability");

  has_M = false;
  m_M = params->get< Teuchos::Array<double> >("Magnetization");
  if( m_M.size()>1 ) has_M = true;

  has_J = false;
  m_J = params->get< Teuchos::Array<double> >("Current Density");
  if( m_J.size()>1 ) has_J = true;
}

// ***********************************************************************
template <typename EvalT>
void SiYuan::EquationSet_Magnetostatics<EvalT>::
buildAndRegisterEquationSetEvaluators(PHX::FieldManager<panzer::Traits>& fm,
                                      const panzer::FieldLibrary& /* fl */,
                                      const Teuchos::ParameterList& /* user_data */) const
{
  using panzer::EvaluatorStyle;

  Teuchos::RCP<panzer::IntegrationRule> ir = this->getIntRuleForDOF(m_dof_name); 
  Teuchos::RCP<panzer::BasisIRLayout> basis = this->getBasisIRLayoutForDOF(m_dof_name); 

  std::string resName("RESIDUAL_" + m_dof_name);

  // Curl Curl Operator
  {
    /*Teuchos::ParameterList p(m_dof_name+"Residual");
    p.set("Residual Name", "RESIDUAL_"+m_dof_name);
    p.set("Value Name", "CURL_"+m_dof_name);
    p.set("Basis", basis);
    p.set("IR", ir);
    const std::vector<std::string> vec = this->getC0();
    p.set("Field Multipliers", Teuchos::rcpFromRef(vec));
    p.set("Multiplier", 1.0);
    
    Teuchos::RCP<PHX::Evaluator<panzer::Traits> > op = Teuchos::rcp(new
      panzer::Integrator_CurlBasisDotVector<EvalT, panzer::Traits>(p));
    fm.template registerEvaluator<EvalT>(op);

    std::string resName = "RESIDUAL_"+m_dof_name;
    std::string valName = "CURL_"+m_dof_name;
    double multiplier(1.0);
    Teuchos::RCP<PHX::Evaluator<panzer::Traits>> op = Teuchos::rcp(new
      panzer::Integrator_CurlBasisDotVector<EvalT, panzer::Traits>(panzer::EvaluatorStyle::CONTRIBUTES,
      resName, valName, *basis, *ir, multiplier));
    this->template registerEvaluator<EvalT>(fm, op);*/

    std::string valName = "CURL_"+m_dof_name;
    Teuchos::RCP<PHX::Evaluator<panzer::Traits>> op = Teuchos::rcp( new
      SiYuan::Residual_Magnetostatics<EvalT, panzer::Traits>(panzer::EvaluatorStyle::EVALUATES,
      resName, valName, *basis, *ir) );
    this->template registerEvaluator<EvalT>(fm, op);
  }
  
  // Source Operator
  if( has_J ) {
    {
      Teuchos::ParameterList pl;
      pl.set("Name","Applied Current");
      pl.set("Data Layout", ir->dl_vector);
      pl.set("Value X", m_J[0]);
      if(basis->getBasis()->dimension()>1)
        pl.set("Value Y", m_J[1]);
      if(basis->getBasis()->dimension()>2)
        pl.set("Value Z", m_J[2]);
    
      Teuchos::RCP<PHX::Evaluator<panzer::Traits> > evaluator  
        = Teuchos::rcp(new panzer::ConstantVector<EvalT,panzer::Traits>(pl));

      fm.registerEvaluator<EvalT>(evaluator);
    }

    std::string valName("Applied Current");
    double multiplier(-1.0);
    Teuchos::RCP<PHX::Evaluator<panzer::Traits>> op = Teuchos::rcp(new
      panzer::Integrator_BasisTimesVector<EvalT, panzer::Traits>(EvaluatorStyle::CONTRIBUTES,
      resName, valName, *basis, *ir, multiplier));
    this->template registerEvaluator<EvalT>(fm, op);
  }

  if( has_M ) {
    {
      Teuchos::ParameterList pl;
      pl.set("Name","Imposed Magnets");
      pl.set("Data Layout", ir->dl_vector);
      pl.set("Value X", m_M[0]);
      if(basis->getBasis()->dimension()>1)
        pl.set("Value Y", m_M[1]);
      if(basis->getBasis()->dimension()>2)
        pl.set("Value Z", m_M[2]);

      Teuchos::RCP<PHX::Evaluator<panzer::Traits> > evaluator  
        = Teuchos::rcp(new panzer::ConstantVector<EvalT,panzer::Traits>(pl));

      fm.registerEvaluator<EvalT>(evaluator);
    }

    std::string valName("Imposed Magnets");
    double multiplier(-1.0);
    Teuchos::RCP<PHX::Evaluator<panzer::Traits>> op = Teuchos::rcp(new
      panzer::Integrator_CurlBasisDotVector<EvalT, panzer::Traits>(EvaluatorStyle::CONTRIBUTES,
      resName, valName, *basis, *ir, multiplier));
    this->template registerEvaluator<EvalT>(fm, op);
  }
}


// ***********************************************************************
/////////////////////////////////////////////////////////////////////////////
  //
  //  Main Constructor: Residual_Magnetostatics
  //
  /////////////////////////////////////////////////////////////////////////////
  template<typename EvalT, typename Traits>
  SiYuan::Residual_Magnetostatics<EvalT, Traits>::
  Residual_Magnetostatics(
    const panzer::EvaluatorStyle&   evalStyle,
    const std::string&              resName,
    const std::string&              valName,
    const panzer::BasisIRLayout&    basis,
    const panzer::IntegrationRule&  ir )
    :
    evalStyle_(evalStyle),
    basisName_(basis.name())
  {
    using PHX::MDField;

    // Ensure the input makes sense.
    TEUCHOS_TEST_FOR_EXCEPTION(resName == "", std::invalid_argument, "Error:  "    \
      "Integrator_CurlBasisDotVector called with an empty residual name.")
    TEUCHOS_TEST_FOR_EXCEPTION(valName == "", std::invalid_argument, "Error:  "    \
      "Integrator_CurlBasisDotVector called with an empty value name.")
    Teuchos::RCP<const panzer::PureBasis> tmpBasis = basis.getBasis();
    TEUCHOS_TEST_FOR_EXCEPTION(not tmpBasis->isVectorBasis(), std::logic_error,
      "Error:  Integrator_CurlBasisDotVector:  Basis of type \""
      << tmpBasis->name() << "\" is not a vector basis.")
    TEUCHOS_TEST_FOR_EXCEPTION(not tmpBasis->requiresOrientations(),
      std::logic_error, "Error:  Integrator_CurlBasisDotVector:  Basis of type \""
      << tmpBasis->name() << "\" does not require orientations, though it "   \
      "should for its use in this Evaluator.")
    TEUCHOS_TEST_FOR_EXCEPTION(not tmpBasis->supportsCurl(), std::logic_error,
      "Error:  Integrator_CurlBasisDotVector:  Basis of type \""
      << tmpBasis->name() << "\" does not support curl.")
    TEUCHOS_TEST_FOR_EXCEPTION(not ((tmpBasis->dimension() == 2) or
      (tmpBasis->dimension() == 3)), std::logic_error,
      "Error:  Integrator_CurlBasisDotVector requires either a two- or "      \
      "three-dimensional basis.  The basis \"" << tmpBasis->name()
      << "\" is neither.")

    // Use a scalar field only if we're dealing with a two-dimensional case.
    spaceDim_ = tmpBasis->dimension();

    // Create the field for the vector-values quantity we're integrating.
    if (spaceDim_ == 2)
    {
      vector2D_ = MDField<const ScalarT, panzer::Cell, panzer::IP>(valName, ir.dl_scalar);
      this->addDependentField(vector2D_);
    }
    else // if (spaceDim_ == 3)
    {
      vector3D_ = MDField<const ScalarT, panzer::Cell, panzer::IP, panzer::Dim>(valName, ir.dl_vector);
      this->addDependentField(vector3D_);
    } // end if spaceDim_ is 2 or 3

    // Create the field that we're either contributing to or evaluating
    // (storing).
    field_ = MDField<ScalarT, panzer::Cell, panzer::BASIS>(resName, basis.functional);
    if (evalStyle == panzer::EvaluatorStyle::CONTRIBUTES)
      this->addContributedField(field_);
    else // if (evalStyle == EvaluatorStyle::EVALUATES)
      this->addEvaluatedField(field_);

    // Add the dependent field multipliers, if there are any.
    std::string mname = "Inverse of Permeability";
    Inv_Permeability_ = MDField<const ScalarT, panzer::Cell, panzer::IP>(mname, ir.dl_scalar);
    this->addDependentField(Inv_Permeability_);

    // Set the name of this object.
    std::string n("Integrator_CurlBasisDotVector (");
    if (evalStyle == panzer::EvaluatorStyle::CONTRIBUTES)
      n += "CONTRIBUTES";
    else // if (evalStyle == EvaluatorStyle::EVALUATES)
      n += "EVALUATES";
    n += "):  " + field_.fieldTag().name();
    this->setName(n);
  } // end of Main Constructor

  /////////////////////////////////////////////////////////////////////////////
  //
  //  postRegistrationSetup()
  //
  /////////////////////////////////////////////////////////////////////////////
  template<typename EvalT, typename Traits>
  void
  SiYuan::Residual_Magnetostatics<EvalT, Traits>:: postRegistrationSetup(
    typename Traits::SetupData sd, PHX::FieldManager<Traits>& fm)
  {
    basisIndex_ = panzer::getBasisIndex(basisName_, (*sd.worksets_)[0], this->wda);

    // Set up the field that will be used to build of the result of this
    // integration.
    /*panzer::MDFieldArrayFactory af("",
      fm.template getKokkosExtendedDataTypeDimensions<EvalT>(), true);
    if (spaceDim_ == 2)
      result2D_ = af.buildStaticArray<ScalarT, panzer::Cell, panzer::IP>(
        "Integrator_CurlBasisDotVector:  2-D Result", vector2D_.extent(0),
        vector2D_.extent(1));
    else // if (spaceDim_ == 3)
      result3D_ = af.buildStaticArray<ScalarT, panzer::Cell, panzer::IP, panzer::Dim>(
        "Integrator_CurlBasisDotVector:  3-D Result", vector3D_.extent(0),
        vector3D_.extent(1), 3);*/
  } // end of postRegistrationSetup()

  /////////////////////////////////////////////////////////////////////////////
  //
  //  evaluateFields()
  //
  /////////////////////////////////////////////////////////////////////////////
  template<typename EvalT, typename Traits>
  void
  SiYuan::Residual_Magnetostatics<EvalT, Traits>::
  evaluateFields( typename Traits::EvalData workset )
  {
    using Kokkos::parallel_for;
    using Kokkos::RangePolicy;
    using PHX::Device;

    typedef Intrepid2::FunctionSpaceTools<PHX::exec_space> FST;

    // Grab the basis information.
    const panzer::BasisValues2<double>& bv = *this->wda(workset).bases[basisIndex_];

    // If we're dealing with a two- or three-dimensional problem...
    if (spaceDim_ == 2)
    {
      PHX::MDField<double, panzer::Cell, panzer::BASIS, panzer::IP>
          weightedCurlBasis = bv.weighted_curl_basis_scalar;
      //FST::integrate<ScalarT>(field_.get_view(), vector2D_.get_view(), weightedCurlBasis.get_view());
      int numBases= weightedCurlBasis.extent(1);
      int numQP = weightedCurlBasis.extent(2);
      for( int cell=0; cell<workset.num_cells; cell++ )
      {
        for (int basis(0); basis < numBases; ++basis)
        {
          if (evalStyle_ == panzer::EvaluatorStyle::EVALUATES)
            field_(cell, basis) = 0.0;
          for (int qp(0); qp < numQP; ++qp) {
            field_(cell, basis) += Inv_Permeability_(cell,qp) * vector2D_(cell, qp) *
                weightedCurlBasis(cell, basis, qp);
          }
        //  std::cout << cell <<"," <<basis << "," << field_(cell, basis) << std::endl;
        } 
      }
    }
    else // if (spaceDim_ == 3)
    {
      PHX::MDField<double, panzer::Cell, panzer::BASIS, panzer::IP, panzer::Dim>
          weightedCurlBasis = bv.weighted_curl_basis_vector;
      //  FST::integrate<ScalarT>(field_.get_view(), vector3D_.get_view(), weightedCurlBasis.get_view());
      int numBases= weightedCurlBasis.extent(1);
      int numQP = weightedCurlBasis.extent(2);
      for( int cell=0; cell<workset.num_cells; cell++ )
      {
        for (int basis(0); basis < numBases; ++basis)
        {
          if (evalStyle_ == panzer::EvaluatorStyle::EVALUATES)
            field_(cell, basis) = 0.0;
          for (int qp(0); qp < numQP; ++qp)
            for (int dim(0); dim<3; ++dim ) {
          //    std::cout << Inv_Permeability_(cell,qp) << std::endl;
              field_(cell, basis) += Inv_Permeability_(cell,qp) * vector3D_(cell, qp, dim) *
                weightedCurlBasis(cell, basis, qp, dim);
            }
        } // end loop over the basis functions
      }
    } // end if spaceDim_ is 2 or 3
  } // end of evaluateFields()

#endif
